package com.yifeng.repo.controller.rocketmq.message;

import com.gomcarter.frameworks.base.common.AssertUtils;
import com.gomcarter.frameworks.base.exception.CustomException;
import com.google.common.base.Strings;
import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.talkyun.utils.revised.Looper;
import com.yifeng.repo.base.api.TimeRange;
import com.yifeng.repo.base.utils.common.BaseUtil;
import com.yifeng.repo.controller.rocketmq.configure.RocketmqProperties;
import com.yifeng.repo.controller.rocketmq.message.consume.AssignPullConsumerLooper;
import lombok.extern.slf4j.Slf4j;
import org.apache.rocketmq.acl.common.AclClientRPCHook;
import org.apache.rocketmq.acl.common.SessionCredentials;
import org.apache.rocketmq.client.consumer.DefaultLitePullConsumer;
import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.common.message.MessageExt;
import org.apache.rocketmq.common.message.MessageQueue;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;
import java.util.stream.Collectors;

/**
 * Created by daibing on 2023/8/5.
 */
@Slf4j
public class PullConsumeAssignManager {
    protected final ConcurrentMap<String, Looper> topicQueue2PullConsumerLooper = new ConcurrentHashMap<>();
    protected final String applicationName;
    protected final RocketmqProperties properties;
    private Cache<String, Collection<MessageQueue>> topic2MessageQueues;

    public PullConsumeAssignManager(String applicationName, RocketmqProperties properties) {
        this.applicationName = applicationName;
        this.properties = properties;
    }

    public void init() {
        topic2MessageQueues = CacheBuilder.newBuilder()
                .initialCapacity(1)
                .maximumSize(8)
                .concurrencyLevel(16)
                .expireAfterWrite(3, TimeUnit.MINUTES)
                .build();
        log.info("yfcloud controller PullConsumeAssignManager init success: OWNER_ID={}, OWNER_NAME={}", BaseUtil.OWNER_ID, BaseUtil.OWNER_NAME);
    }

    public void destroy() {
        topic2MessageQueues.cleanUp();
        topicQueue2PullConsumerLooper.forEach((key, value) -> value.shutdown());
        log.info("yfcloud controller PullConsumeAssignManager destroy success: OWNER_ID={}, OWNER_NAME={}", BaseUtil.OWNER_ID, BaseUtil.OWNER_NAME);
    }

    /**
     * 增加Topic队列单独消费（使用分配模式：AssignPullConsumerLooper）
     *
     * @param topic           消息主题
     * @param queueId         队列id
     * @param msgListConsumer 拉取消息的消费者
     * @param timeRangeList   拉取消息的时间段
     */
    public void addTopicQueueConsumer(String topic, int queueId, List<TimeRange> timeRangeList, Consumer<List<MessageExt>> msgListConsumer) {
        addTopicQueueConsumer(topic, queueId, null, timeRangeList, msgListConsumer);
    }

    /**
     * 增加Topic队列单独消费（使用分配模式：AssignPullConsumerLooper），在上一个方法基础上支持多消费组
     *
     * @param timeRangeCode 指定队列使用指定时间区间编号对应的消费组，如果没有指定时间区间编号，或者如果没有配置，都使用默认消费组
     */
    public void addTopicQueueConsumer(String topic, int queueId, String timeRangeCode, List<TimeRange> timeRangeList, Consumer<List<MessageExt>> msgListConsumer) {
        String key = topic + "-queueId-" + queueId;
        Looper looper = topicQueue2PullConsumerLooper.get(key);
        if (looper != null) {
            return;
        }
        synchronized (topicQueue2PullConsumerLooper) {
            looper = topicQueue2PullConsumerLooper.get(key);
            if (looper != null) {
                return;
            }
            looper = buildAssignPullConsumerLooper(topic, queueId, timeRangeCode, timeRangeList, msgListConsumer);
            looper.start();
            topicQueue2PullConsumerLooper.put(key, looper);
        }
    }

    /**
     * 移除Topic队列单独消费
     */
    public void removeTopicQueueConsumer(String topic, int queueId) {
        String key = topic + "-queueId-" + queueId;
        Looper looper = topicQueue2PullConsumerLooper.remove(key);
        if (looper != null) {
            looper.shutdown();
            log.info("removeTopicQueueConsumer success: {}-{}", topic, queueId);
        }
    }

    /**
     * 列出当前消息主题的消费queueIds（分配模式）
     */
    public List<Integer> listConsumerQueueId(String topic) {
        return topicQueue2PullConsumerLooper.keySet().stream()
                .filter(k -> k.startsWith(topic + "-queueId-"))
                .map(s -> Integer.parseInt(s.substring((topic + "-queueId-").length())))
                .collect(Collectors.toList());
    }

    /**
     * 查询指定topic的消息队列
     */
    public List<MessageQueue> listMessageQueue(String topic) {
        DefaultLitePullConsumer pullConsumer = getConsumerFromLooper(topic);
        try {
            return new ArrayList<>(pullConsumer.fetchMessageQueues(topic));
        } catch (MQClientException e) {
            throw new CustomException(e);
        }
    }

    /**
     * 获取队列的统计数据
     *
     * @param topic     消息主题
     * @param timestamp 开始时间，单位毫秒
     * @return queueId -> searchOffset
     */
    public Map<Integer, Long> listMessageQueueStats(String topic, long timestamp) {
        DefaultLitePullConsumer pullConsumer = getConsumerFromLooper(topic);
        Map<Integer, Long> mqStats = new HashMap<>();
        try {
            Collection<MessageQueue> messageQueues = pullConsumer.fetchMessageQueues(topic);
            for (MessageQueue messageQueue : messageQueues) {
                mqStats.put(messageQueue.getQueueId(), pullConsumer.offsetForTimestamp(messageQueue, timestamp));
            }
        } catch (MQClientException e) {
            throw new CustomException(e);
        }
        return mqStats;
    }

    protected DefaultLitePullConsumer getConsumerFromLooper(String topic) {
        AssertUtils.notEmpty(topicQueue2PullConsumerLooper.values(), "当前还没有初始化消费者，不能获取消息队列列表！");
        Looper looper = topicQueue2PullConsumerLooper.entrySet()
                .stream()
                .filter(entry -> entry.getKey().startsWith(topic + "-"))
                .findFirst()
                .orElseThrow(() -> new CustomException("获取不到Topic对应的消费者：" + topic)).getValue();
        return ((AssignPullConsumerLooper) looper).getPullConsumer();
    }

    protected Looper buildAssignPullConsumerLooper(String topic, int queueId, String timeRangeCode, List<TimeRange> timeRangeList,
                                                   Consumer<List<MessageExt>> msgListConsumer) {
        String key = topic + "-queueId-" + queueId;
        DefaultLitePullConsumer pullConsumer = buildAssignPullConsumer(topic, queueId, timeRangeCode);
        MessageQueue messageQueue = getMessageQueueByCache(topic, queueId, pullConsumer);
        RocketmqProperties.TopicProperties topicProperties = getTopicProperties(properties, topic);
        return new AssignPullConsumerLooper(key + "-looper", 0, 1000, messageQueue,
                pullConsumer, timeRangeList, msgListConsumer, topicProperties.isSeekBackWhenError());
    }

    /**
     * 分配模式：当前一个队列创建一个consumer
     */
    protected DefaultLitePullConsumer buildAssignPullConsumer(String topic, int queueId, String timeRangeCode) {
        // 客户端实例名称
        RocketmqProperties.TopicProperties topicProperties = getTopicProperties(properties, topic);

        // 指定队列使用指定时间区间编号对应的消费组，如果没有指定时间区间编号，或者如果没有配置，都使用默认消费组
        String consumerGroup = Strings.isNullOrEmpty(timeRangeCode) || topicProperties.getConsumerGroupMapping() == null
                ? topicProperties.getConsumerGroup()
                : topicProperties.getConsumerGroupMapping().getOrDefault(timeRangeCode, topicProperties.getConsumerGroup());
        AclClientRPCHook rpcHook = new AclClientRPCHook(new SessionCredentials(topicProperties.getAccessKey(), topicProperties.getSecretKey()));

        // build pull consumer
        DefaultLitePullConsumer pullConsumer = new DefaultLitePullConsumer(consumerGroup, rpcHook);
        pullConsumer.setInstanceName(applicationName + "-" + topic + "-pull-assign-" + queueId + "-" + consumerGroup);
        pullConsumer.setNamesrvAddr(topicProperties.getNamesrvAddr());
        pullConsumer.setAutoCommit(false);
        pullConsumer.setPullBatchSize(topicProperties.getPullBatchSize());
        pullConsumer.setPullThreadNums(topicProperties.getAssignPullThreadNums());
        try {
            pullConsumer.start();
            MessageQueue messageQueue = getMessageQueueByCache(topic, queueId, pullConsumer);
            pullConsumer.assign(Collections.singleton(messageQueue));
            log.info("buildAssignPullConsumer success: {}-{}, consumerGroup={}, pullBatchSize={}, pullThreadNums={}",
                    topic, queueId, pullConsumer.getConsumerGroup(), pullConsumer.getPullBatchSize(), pullConsumer.getPullThreadNums());
            return pullConsumer;
        } catch (MQClientException e) {
            throw new CustomException("buildAssignPullConsumer failed " + topic + "-" + queueId + "-" + consumerGroup, e);
        }
    }

    protected MessageQueue getMessageQueueByCache(String topic, int queueId, DefaultLitePullConsumer pullConsumer) {
        try {
            Collection<MessageQueue> messageQueues = topic2MessageQueues.get(topic, () -> pullConsumer.fetchMessageQueues(topic));
            return messageQueues.stream()
                    .filter(q -> q.getQueueId() == queueId)
                    .findFirst()
                    .orElseThrow(() -> new CustomException(topic + "没有对应的Queue：" + queueId));
        } catch (ExecutionException e) {
            throw new CustomException(e);
        }
    }

    protected RocketmqProperties.TopicProperties getTopicProperties(RocketmqProperties properties, String topic) {
        return properties.getTopicPropertiesList()
                .stream()
                .filter(s -> s.getTopic().equals(topic))
                .findFirst()
                .orElseThrow(() -> new CustomException("没有Topic配置：" + topic));
    }

}
